package pl;

import javacard.framework.ISO7816;
import javacard.framework.ISOException;
import javacard.framework.JCSystem;
import javacard.framework.Util;
import javacard.security.MessageDigest;
import javacard.security.RandomData;

public class EMSA_PSS {

	private static final byte _0X_BC = (byte) 0xBC;

	private static final byte[] ZERO_ARRAY = new byte[8];

	protected MessageDigest md;
	protected MessageDigest mgfDig;
	private RandomData rd;
	
	public short hLen;
	public byte[] hashResult;

	byte[] salt;
	byte[] result;
	byte[] dbMask;
	byte[] DB;
	private byte[] C;
	byte[] hash;

	short emLen;
    short saltLen;
    
    private short emBits;
	private short counter;
	private short hLen2;
	private short dbMaskLen;
	
	public EMSA_PSS() {

		// default salt length for PSS
		this.saltLen = 20;
		
		// default - sha1
		md = MessageDigest.getInstance(MessageDigest.ALG_SHA, false);

		hLen = (short) md.getLength();
		
		hashResult = JCSystem.makeTransientByteArray(hLen,
				JCSystem.CLEAR_ON_DESELECT);

		hash = JCSystem
				.makeTransientByteArray(hLen, JCSystem.CLEAR_ON_DESELECT);

		rd = RandomData.getInstance(RandomData.ALG_PSEUDO_RANDOM);

		salt = JCSystem
				.makeTransientByteArray(saltLen, JCSystem.CLEAR_ON_DESELECT);

		C = JCSystem.makeTransientByteArray((short) 4,
				JCSystem.CLEAR_ON_DESELECT);

		
		result = JCSystem.makeTransientByteArray(MediatorEngine.MAX_EMLEN,
				JCSystem.CLEAR_ON_DESELECT);

		DB = JCSystem.makeTransientByteArray(MediatorEngine.MAX_EMLEN,
				JCSystem.CLEAR_ON_DESELECT);

		// default for mgf1
		mgfDig = MessageDigest.getInstance(MessageDigest.ALG_SHA, false);
		
	}

	public void init(short saltLen, short modulusLen)
	{
		this.saltLen = saltLen;
		emBits = (short) (modulusLen - 1);
		emLen = (short) ((short) (emBits + 7) / 8);	
		dbMaskLen = (short) (emLen - hLen - 1);
	}
	
	public void emsaPssEncode(byte[] hashedMessage) {

		rd.generateData(salt, (short) 0, saltLen);
		
		if (emBits < (short) (8 * hLen + 8 * saltLen + 9)) {
			/* TODO: Exception */
			//ISOException.throwIt((short) (ISO7816.SW_WRONG_LENGTH + emBits));
			
			ISOException.throwIt(ISO7816.SW_WRONG_LENGTH);
		}

		// hash (00 00 00 00 00 00 00 00 || hash || salt)
		md.reset();
		md.update(ZERO_ARRAY, (short) 0, (short) ZERO_ARRAY.length);
		md.update(hashedMessage, (short) 0, (short) hashedMessage.length);

		// add salt
		md.doFinal(salt, (short) 0, (short) saltLen, hashResult, (short) 0);
		
		//Utils.ClearArray(DB);
		Utils.ClearArray(DB, (short) (emLen - hLen - 1));

		DB[(short) (emLen - saltLen - hLen - 2)] = 0x01;

		Util.arrayCopy(salt, (short) 0, DB, (short) (emLen - saltLen - hLen - 1),
				saltLen);

		// mgf1
		// dbMask = null;
		// dbMask = new byte[(short) (emLen - hLen - 1)];
		// dbMask = mgf1(hashResult, (short) (emLen - hLen - 1));
		mgf1(hashResult, (short)0, hLen, result, (short)0);

		for (short i = 0; i < dbMaskLen; i++)
			DB[i] = (byte) (DB[i] ^ result[i]);
		
		DB[0] &= (0xFF >>> (8 * emLen - emBits));

		Utils.ClearArray(result, (short) (emLen - hLen - 1));
		
		Util.arrayCopy(DB, (short) 0, result, (short) 0,
				(short) (emLen - hLen - 1));
		Util.arrayCopy(hashResult, (short) 0, result,
				(short) (emLen - hLen - 1), hLen);

		result[(short) (emLen - 1)] = _0X_BC;
	}
	
	public short emsaPssEncode(byte[] message, short messageOffset, short messageLen, byte[] output) {

		
		rd.generateData(salt, (short) 0, saltLen);
		
		if (emBits < (short) (8 * messageLen + 8 * saltLen + 9)) {
			/* TODO: Exception */
			//ISOException.throwIt((short) (ISO7816.SW_WRONG_LENGTH + emBits));
			
			ISOException.throwIt(ISO7816.SW_WRONG_LENGTH);
		}

		// hash (00 00 00 00 00 00 00 00 || hash || salt)
		md.reset();
		md.update(ZERO_ARRAY, (short) 0, (short) ZERO_ARRAY.length);
		md.update(message, messageOffset, messageLen);

		// add salt
		md.doFinal(salt, (short) 0, (short) saltLen, hashResult, (short) 0);
		
		//Utils.ClearArray(DB);
		Utils.ClearArray(DB, (short) (emLen - hLen - 1));

		DB[(short) (emLen - saltLen - hLen - 2)] = 0x01;

		Util.arrayCopy(salt, (short) 0, DB, (short) (emLen - saltLen - hLen - 1),
				saltLen);

		// mgf1
		// dbMask = null;
		// dbMask = new byte[(short) (emLen - hLen - 1)];
		// dbMask = mgf1(hashResult, (short) (emLen - hLen - 1));
		mgf1(hashResult, (short)0, hLen, result, (short)0);

		for (short i = 0; i < dbMaskLen; i++)
			DB[i] = (byte) (DB[i] ^ result[i]);
		
		DB[0] &= (0xFF >>> (8 * emLen - emBits));

		Utils.ClearArray(result, (short) (emLen - hLen - 1));
		
		Util.arrayCopy(DB, (short) 0, result, (short) 0,
				(short) (emLen - hLen - 1));
		Util.arrayCopy(hashResult, (short) 0, result,
				(short) (emLen - hLen - 1), hLen);

		result[(short) (emLen - 1)] = _0X_BC;
		
		Util.arrayCopy(result, (short)0, output, (short)0, emLen);
		
		return emLen;
	}

	public boolean emsaPssVerify(byte[] message, short messageLen, byte[] eM) {

		md.reset();
		md.doFinal(message, (short)0, messageLen, this.hashResult, (short)0);
		
		// emLen
		if ((byte) (eM[(short) (emLen - 1)]) != (_0X_BC)) {
			return false;
		}

		// out buff -> 256 array

		Utils.ClearArray(DB);

		mgf1(eM, (short) (emLen - hLen - 1), hLen, DB,
				(short) 0);

		// ClearArray(result);

		for (short i = 0; i != dbMaskLen; i++) {
			eM[i] ^= DB[i];
		}

		eM[0] &= (0xff >> ((emLen * 8) - emBits));

		for (short i = 0; i != (short) (emLen - hLen - saltLen - 2); i++) {
			if (eM[i] != 0) {
				// clearBlock(block);
				return false;
			}
		}

		if (eM[(short) (emLen - hLen - saltLen - 2)] != 0x01) {
			// clearBlock(block);
			return false;
		}
		//**********************************************************//*
		
		short saltOff = (short) (emLen - saltLen - hLen - 1); // in DB

		md.reset();
		md.update(ZERO_ARRAY, (short) 0, (short) ZERO_ARRAY.length);
		md.update(hashResult, (short) 0, (short) hLen);
		
		md.doFinal(eM, saltOff, saltLen, hashResult, (short)0);
		
		for (short i = (short) (emLen - hLen - 1), j = (short) 0; j != hLen; i++, j++) {
			if ((eM[i] ^ hashResult[j]) != 0) {
				return false;
			}
		}

		return true;

	}

	private void mgf1(byte[] seed, short offset, short seedLen, byte[] outBuffer, short cnt) {

		Utils.ClearArray(outBuffer);
		Utils.ClearArray(hash);

		counter = cnt;
		hLen2 = mgfDig.getLength();

		mgfDig.reset();

		while (counter < (short) ((short) dbMaskLen / hLen2)) {
			I2OSP(counter, C);
			mgfDig.update(seed, offset, seedLen);
			mgfDig.doFinal(C, (short) 0, (short) C.length, hash, (short) 0);
			
			Util.arrayCopy(hash, (short) 0, outBuffer, (short) (counter * hLen2),
					hLen2);
			counter++;
		}

		if ((short) (counter * hLen2) < (short) dbMaskLen) {
			I2OSP(counter, C);
			mgfDig.update(seed, offset, seedLen);
			mgfDig.doFinal(C, (short) 0, (short) C.length, hash, (short) 0);
			Util.arrayCopy(hash, (short) 0, outBuffer, (short) (counter * hLen2),
					(short) (dbMaskLen - (short) (counter * hLen2)));
		}
	}
	
	public static void I2OSP(short i, byte[] sp) {
		sp[0] = (byte) (i >>> 24);
		sp[1] = (byte) (i >>> 16);
		sp[2] = (byte) (i >>> 8);
		sp[3] = (byte) (i >>> 0);
	}
	
	public void emsaPkcsEncode(byte[] message, short messageLen)
	{
		//  EB = 00 || 01 || min 8 x FF || 00 || message
		hLen = messageLen;
		
		if (emBits < (short) (8 * hLen + 8 * saltLen + 9)) {
			/* TODO: Exception */
			//ISOException.throwIt((short) (ISO7816.SW_WRONG_LENGTH + emBits));
			
			ISOException.throwIt(ISO7816.SW_WRONG_LENGTH);
		}
		
		
	}
}